''' 3d morphable model example
3dmm parameters --> mesh 
fitting: 2d image + 3dmm -> 3d face
'''
import os, sys
import subprocess
import numpy as np
import scipy.io as sio
from skimage import io
from time import time
import matplotlib.pyplot as plt
import re
import math

sys.path.append('..')
import face3d
from face3d import mesh
from face3d.morphable_model import MorphabelModel
from face3d import mesh
from face3d.morphable_model.fit_multi import MultiFit

import cv2
import dlib
import string

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('../Data/shape_predictor_68_face_landmarks.dat')

# --------------------- Forward: parameters(shape, expression, pose) --> 3D obj --> 2D image  ---------------
# --- 1. load model
bfm = MorphabelModel('../Data/BFM.mat')
print('init bfm model success')

# --- 2. generate face mesh: vertices(represent shape) & colors(represent texture)
sp = bfm.get_shape_para('random')
ep = bfm.get_exp_para('random')
vertices = bfm.generate_vertices(sp, ep)

tp = bfm.get_tex_para('random')
colors = bfm.generate_colors(tp)
colors = np.minimum(np.maximum(colors, 0), 1)

# --- 3. transform vertices to proper position
s = 10e-04              # 比例
angles = [0, -30, 0]    # 观察角度
t = [0, 0, 0]           # 平移量
transformed_vertices = bfm.transform(vertices, s, angles, t)
projected_vertices = transformed_vertices.copy() # using stantard camera & orth projection

# --- 4. render(3d obj --> 2d image)
# set prop of rendering
h = w = 256     # 输出图像宽高
c = 3           #
image_vertices = mesh.transform.to_image(projected_vertices, h, w)
image = mesh.render.render_colors(image_vertices, bfm.triangles, colors, h, w)

# -------------------- Back:  2D image points and corresponding 3D vertex indices-->  parameters(pose, shape, expression) ------
## only use 68 key points to fit
x = projected_vertices[bfm.kpt_ind, :2] # 2d keypoint, which can be detected from image
#num = 1
#with open("e:/8d/test_landmarks.txt", "rt") as fp:
#    num = fp.readline()
#    num = int(num)
#    #fp.write("{0}\n".format(shape.num_parts))
#    for i in range(num):
#        #fp.write("{0} {1}\n".format(shape.part(i).x, shape.part(i).y))
#        xs = fp.readline()
#        b = re.findall(r"\d+\.?\d*",xs)
#        #b= [int(i) for i in xs if xs.isdigit()]
#        #x[i][0] = int(b[0])
#        #x[i][1] = int(b[1])
X_ind = bfm.kpt_ind # index of keypoints in 3DMM. fixed.



#image_vertices = mesh.transform.to_image(transformed_vertices, h, w)
#fitted_image = mesh.render.render_colors(image_vertices, bfm.triangles, colors, h, w)

#save_folder = 'results/3dmm'
#if not os.path.exists(save_folder):
#    os.mkdir(save_folder)

#io.imsave('{}/generated.jpg'.format(save_folder), image)
#io.imsave('{}/fitted.jpg'.format(save_folder), fitted_image)


'''
xp2 = xp1 - x
xmax = xp2.max()
xmin = xp2.min()
np.reshape(xp2, (68*2,))
xp3 = xp2.flatten();
xp4 = xp3.dot(xp3)
print(math.sqrt(xp4/68), xmin, xmax)
# ------------- print & show 
print('pose, groudtruth: \n', s, angles[0], angles[1], angles[2], t[0], t[1])
print('pose, fitted: \n', fitted_s0, fitted_angles0[0], fitted_angles0[1], fitted_angles0[2], fitted_t0[0], fitted_t0[1])

### ----------------- visualize fitting process
# fit
fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t = bfm.fit(x, X_ind, max_iter = 3, isShow = True)

print("default:", s, angles, t);

# 显示解的过程
# verify fitted parameters
for i in range(fitted_sp.shape[0]):
    fitted_vertices = bfm.generate_vertices(fitted_sp[i], fitted_ep[i])
    #transformed_vertices = bfm.transform(fitted_vertices, fitted_s[i], fitted_angles[i], fitted_t[i])
    transformed_vertices = bfm.transform(fitted_vertices, s, angles, t)
    print(i, fitted_s[i], fitted_angles[i], fitted_t[i]);

    image_vertices = mesh.transform.to_image(transformed_vertices, h, w)
    fitted_image = mesh.render.render_colors(image_vertices, bfm.triangles, colors, h, w)
    io.imsave('{}/show_{:0>2d}.jpg'.format(save_folder, i), fitted_image)
'''

#显示29个表情的变化
## verify fitted parameters
##for i in range(fitted_sp.shape[0]):
#i = 0
#for j in range(29):
#    bak = fitted_ep[0][j][0]
#    fitted_ep[0][j][0]=-2
#    fitted_vertices = bfm.generate_vertices(fitted_sp[i], fitted_ep[i])
#    #transformed_vertices = bfm.transform(fitted_vertices, fitted_s[i], fitted_angles[i], fitted_t[i])
#    transformed_vertices = bfm.transform(fitted_vertices, s, angles, t)

#    image_vertices = mesh.transform.to_image(transformed_vertices, h, w)
#    fitted_image = mesh.render.render_colors(image_vertices, bfm.triangles, colors, h, w)
#    io.imsave('{}/show_{:0>2d}.jpg'.format(save_folder, j), fitted_image)
#    fitted_ep[0][j][0] = bak

options = '-delay 20 -loop 0 -layers optimize' # gif. need ImageMagick.
subprocess.call('convert {} {}/show_*.jpg {}'.format(options, save_folder, save_folder + '/3dmm.gif'), shell=True)
subprocess.call('rm {}/show_*.jpg'.format(save_folder), shell=True)
